import numpy as np
import socket
import json
from typing import Dict, Optional


# 观测信息类定义
class EnemyInfo:
    """敌军信息类"""

    def __init__(self, raw_data):
        self.EnemyID = None  # 敌机ID
        self.isNTS = False  # 被火控锁定，isNTS=1，否则为0
        self.TargetDis = 0  # 目标距离 m
        self.DisRate = 0  # 目标距离变化率 m/s
        self.TargetYaw = 0  # 目标方位角 rad
        self.TargetPitch = 0  # 目标俯仰角 rad
        self.TargetV_N = 0  # 目标北向速度 m/s
        self.TargetV_U = 0  # 目标天向速度
        self.TargetV_E = 0  # 目标东向速度
        self.TargetMach_M = 0  # 目标马赫数M
        self.TargetAlt = 0  # 目标高度


class AttackInfo:
    """攻击信息类"""

    def __init__(self, raw_data):
        self.EnemyID = None  # 敌机ID
        self.NTSstate = 0  # 火控状态，即目标飞机是否被火控锁定的标志位状态，只有开启和关闭两个状态
        # NTSstate=1：敌机被雷达发现但未被火控锁定，NTSstate=2：被我方火控锁定
        self.TargetDis = 0  # 目标距离 m
        self.DisRate = 0  # 目标距离变化率 m/s
        self.MissilePowerfulDis = 0  # 有较大把握打中的距离 km
        self.MissileMaxDis = 0  # 有可能打中的最远距离 km


class MissileOutInfo:
    """导弹发射信息类"""

    def __init__(self, raw_data_1, raw_data_2):
        self.RawDataProperty = raw_data_1  # MslProperty_data
        self.RawDataInfoBack = raw_data_2  # AIMInfoBack_data
        self.WeaponID = None  # 武器编号
        self.state = 0  # 发射后导引头状态（0=导弹未到达导引头开机点，1=导弹到达导引头开机点 2=导弹导引头已截获，3=导弹被动头已经截获）
        self.TargetID = None  # 该武器对应目标
        self.Yaw = 0  # 导弹航向 rad
        self.Pitch = 0  # 导弹俯仰角 rad
        self.fromDroneDis = 0  # 导弹与载机距离 km
        self.openSeeker = 0  # 导弹导引头开机点位置
        self.isTracking = False  # 导弹截获目标标志（0=N，1=Y）（末制导跟踪目标）
        self.TargetDis = 0  # 导弹-目标相对距离 km
        self.Vx = 0  # 导弹速度各向分量 m/s
        self.Vy = 0
        self.Vz = 0
        self.x = 0  # 导弹位置坐标(半稳系) km
        self.y = 0
        self.z = 0


class AlarmInfo:
    """告警信息类"""

    def __init__(self, raw_data):
        self.AlarmID = None  # 告警信息辐射源ID
        self.AlarmType = "导弹"  # 告警类别:导弹、飞机
        self.MisAzi = 0  # 相对方位rad


class DroneInfo:
    """无人机信息类"""

    def __init__(self, raw_data):
        # 无人机机载信息
        self.DroneID = None  # 本机ID
        self.Longtitude = 0  # 经度rad
        self.Latitude = 0  # 纬度rad
        self.Altitude = 0  # 海拔高度m
        self.Yaw = 0  # 载机航向角(右偏为正) 单位：rad
        self.Pitch = 0  # 载机俯仰角 单位：rad
        self.Roll = 0  # 载机滚转角 单位：rad
        self.V_N = 0  # 北向速度 m/s
        self.V_U = 0  # 天向速度
        self.V_E = 0  # 东向速度
        self.A_N = 0  # 北向加速度 m/s**2
        self.A_U = 0  # 天向加速度
        self.A_E = 0  # 东向加速度
        self.Mach_M = 0  # 马赫速 M
        self.NormalAccG = 0  # 法向过载 g
        # 传感器信息
        self.isRadiateOn = True  # 雷达辐射状态
        self.ScanScale = 0  # 雷达扫描宽度
        self.isMisWarning = False  # 导弹逼近告警
        self.AlarmList = []  # 告警信息列表
        # 武器信息
        self.MissileNowNum = 0  # 现有导弹数量
        self.AttackPowerfulDis = 0  # 发射后可导弹自己制导距离
        self.AttackAllowDis = 0  # 最大可发射距离
        self.MissileTrackList = []  # 本机已发射导弹
        # 敌方信息
        self.FoundEnemyList = []  # 本机发现敌方列表
        self.AttackEnemyList = []  # 本机攻击的目标列表


# 控制指令类定义
class OtherControlInfo:
    """其他控制信息类"""

    def __init__(self):
        self.isLaunch = 0  # 用于控制导弹的发射，要求上升沿有效


class PlaneControlInfo:
    """飞行控制信息类"""

    def __init__(self):
        # 推杆控制方式参数
        self.ForceLat = 0  # 滚转指令   操纵杆横向指令（右倾为+,范围-80~+80）
        self.ForcePedal = 0  # 偏航指令   脚蹬（左脚蹬为+,范围-500~+500）
        self.ForceLong = 0  # 俯仰指令   操纵杆纵向指令（后拉为+，范围-80~+160）
        self.ThrottlePosition = 0  # 油门杆位置 0~120

        # 机动号控制方式参数
        self.CmdIndex = 1  # 动作序号 从1开始累加 参数变化而机动类型不变时可以不更新 建议按照机动指令的顺序设置
        self.CmdID = 1  # 动作编号
        self.isApplyNow = False  # true:立即应用 false：完成前序命令后应用 目前默认为true
        self.VelType = 2  # V_type 目标速度类型  0 马赫数 1 真空速(km/h) 2 表速(m/s)
        self.CmdNy = 1  # 法向过载 单位：g  1~8
        self.CmdSpd = 0  # 目标速度或限制速度 m/s 或 mach
        self.CmdAlt = 1000  # safe 安全高度 目标高度 m 与飞行包线有关 <16000m
        self.TurnDirection = 0  # 转弯方向 1 右转 -1左转 0默认就近转
        self.CmdDeltaHeading = 0  # 最大水平偏置角 度-180~+180  S形机动的航向偏置角
        self.Time = 0  # 时间 秒 S机动直线段时间   S形机动的直线段飞行时间
        self.CmdThrust = 0  # 推力指令 0~120
        self.ThrustLimit = 100  # 最大推力限制 0~120 可默认为100
        self.CmdPitchDeg = 0  # 目标航迹倾斜角 度 -90~+90
        self.CmdHeadingDeg = 0  # 目标航向 度 -180~+180
        self.CmdObliqueDeg = 0  # 向上/向下斜筋斗轨迹平面倾角
        self.CmdRollDeg = 0  # 目标姿态滚转角度 -180~+180   目标滚转角


class SOCtrlInfo:
    """控制接口类"""

    def __init__(self):
        self.NTSEntityIdAssigned = 1  # 飞机火控锁定的目标实体ID
        self.isNTSAssigned = 0  # 用于控制火控锁定目标ID的修改，上升沿使锁定目标改为NTSEntityIdAssigned


class AirCombatEnv:
    """
    空战环境类
    观测空间: 64 * 64 * 5 (3维)
    动作空间: 64维向量 (1 + 9个方向 * (6种移动 + 1种攻击))
    """

    def __init__(self, args):

        self.args = args
        self.n_agents = args.n_agents
        self.n_enemies = args.n_enemies
        self.episode_limit = args.episode_limit

        # 地图参数
        self.map_lon_min = args.map_lon_min
        self.map_lon_max = args.map_lon_max
        self.map_lat_min = args.map_lat_min
        self.map_lat_max = args.map_lat_max

        # 威胁区域参数（半球形威胁区）
        self.threat_center_lon = args.threat_center_lon  # 威胁区中心经度
        self.threat_center_lat = args.threat_center_lat  # 威胁区中心纬度
        self.threat_center_alt = args.threat_center_alt  # 威胁区中心高度
        self.threat_damage_radius = args.threat_damage_radius  # 威胁区毁伤有效半径
        self.earth_radius = 6371000.0  # 地球半径（米）

        # 海拔高度参数
        self.altitude_min = args.altitude_min
        self.altitude_max = args.altitude_max

        # 观测空间和动作空间
        self.obs_shape = (36,)  # 36维状态向量
        self.act_shape = 7  # 动作空间维度: 7个动作

        # 网络连接
        self.socket = None
        self.connected = False

        # 环境观测
        self.step_count = 0
        self.episode_reward = 0
        self.done = False

        # 智能体观测
        self.agents_info = []  # 存储所有智能体的DroneInfo
        self.enemies_info = []  # 存储所有敌机的信息

        # 初始化网络连接
        self._connect_to_simulation()

    def _connect_to_simulation(self):
        """连接到仿真平台"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect(('localhost', 8080))  # 根据实际IP和端口修改
            self.connected = True
            print("成功连接到仿真平台")
        except Exception as e:
            print(f"连接仿真平台失败: {e}")
            self.connected = False

    def _disconnect_from_simulation(self):
        """断开与仿真平台的连接"""
        if self.socket:
            self.socket.close()
            self.connected = False
            print("已断开与仿真平台的连接")

    def _calculate_threat_zone_ratio_in_cell(self, grid_x: int, grid_y: int, grid_h: int) -> float:
        """
        计算每个小立体空间内半球威胁区的体积占比
        
        Args:
            grid_x: 网格X坐标 (0-31)
            grid_y: 网格Y坐标 (0-31) 
            grid_h: 网格高度层 (0-2)
            
        Returns:
            float: 该小立体空间内威胁区体积占比 [0.0, 1.0]
        """
        import math

        # 计算每个小立体空间的尺寸
        lon_range = self.map_lon_max - self.map_lon_min
        lat_range = self.map_lat_max - self.map_lat_min
        alt_range = self.altitude_max - self.altitude_min

        # 每个小立体空间的经纬度范围
        cell_lon_min = self.map_lon_min + (grid_x / 32.0) * lon_range
        cell_lon_max = self.map_lon_min + ((grid_x + 1) / 32.0) * lon_range
        cell_lat_min = self.map_lat_min + (grid_y / 32.0) * lat_range
        cell_lat_max = self.map_lat_min + ((grid_y + 1) / 32.0) * lat_range
        cell_alt_min = self.altitude_min + (grid_h / 3.0) * alt_range
        cell_alt_max = self.altitude_min + ((grid_h + 1) / 3.0) * alt_range

        # 将经纬度范围转换为米
        earth_radius = 6371000.0  # 地球半径（米）
        center_lat_rad = math.radians((cell_lat_min + cell_lat_max) / 2)

        # 小立体空间的尺寸（米）
        cell_lon_distance = earth_radius * math.cos(center_lat_rad) * math.radians(cell_lon_max - cell_lon_min)
        cell_lat_distance = earth_radius * math.radians(cell_lat_max - cell_lat_min)
        cell_alt_distance = cell_alt_max - cell_alt_min

        # 小立体空间的体积
        cell_volume = cell_lon_distance * cell_lat_distance * cell_alt_distance

        # 计算半球威胁区在小立体空间内的体积
        # 使用数值积分方法近似计算
        threat_volume_in_cell = self._calculate_hemisphere_volume_in_cell(
            cell_lon_min, cell_lon_max, cell_lat_min, cell_lat_max,
            cell_alt_min, cell_alt_max, cell_lon_distance, cell_lat_distance, cell_alt_distance
        )

        # 计算占比
        if cell_volume > 0:
            volume_ratio = threat_volume_in_cell / cell_volume
            return min(1.0, volume_ratio)
        else:
            return 0.0

    def _calculate_hemisphere_volume_in_cell(self, lon_min, lon_max, lat_min, lat_max,
                                             alt_min, alt_max, lon_dist, lat_dist, alt_dist):
        """
        计算半球威胁区在小立体空间内的体积（数值积分）
        """
        import math

        # 使用蒙特卡洛方法近似计算
        sample_points = 1000  # 采样点数量
        points_in_hemisphere = 0

        for _ in range(sample_points):
            # 在小立体空间内随机采样
            sample_lon = lon_min + (lon_max - lon_min) * np.random.random()
            sample_lat = lat_min + (lat_max - lat_min) * np.random.random()
            sample_alt = alt_min + (alt_max - alt_min) * np.random.random()

            # 计算采样点到威胁区中心的距离
            lon1_rad = math.radians(self.threat_center_lon)
            lat1_rad = math.radians(self.threat_center_lat)
            lon2_rad = math.radians(sample_lon)
            lat2_rad = math.radians(sample_lat)

            # 计算球面距离
            dlat = lat2_rad - lat1_rad
            dlon = lon2_rad - lon1_rad
            a = (math.sin(dlat / 2) ** 2 +
                 math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(dlon / 2) ** 2)
            c = 2 * math.asin(math.sqrt(a))
            horizontal_distance = self.earth_radius * c

            # 计算3D距离
            height_diff = sample_alt - self.threat_center_alt
            distance_3d = math.sqrt(horizontal_distance ** 2 + height_diff ** 2)

            # 判断是否在半球威胁区内
            if distance_3d <= self.threat_damage_radius:
                points_in_hemisphere += 1

        # 计算体积占比
        volume_ratio = points_in_hemisphere / sample_points
        cell_volume = lon_dist * lat_dist * alt_dist

        return volume_ratio * cell_volume

    def _send_control_command(self, agent_id: int, control_info: Dict):
        """发送控制指令"""
        if not self.connected:
            return False

        try:
            command = {
                'agent_id': agent_id,
                'control': control_info
            }
            self.socket.send(json.dumps(command).encode())
            return True
        except Exception as e:
            print(f"发送控制指令失败: {e}")
            return False

    def _receive_obs_info(self) -> Dict:
        """接收观测信息"""
        if not self.connected:
            return {}

        try:
            data = self.socket.recv(4096)
            if data:
                return json.loads(data.decode())
            return {}
        except Exception as e:
            print(f"接收观测信息失败: {e}")
            return {}

    def _create_4d_obs_map(self, agent_id: int) -> np.ndarray:
        """
        创建4D观测地图 (32 * 32 * 3 * 5)
        
        Args:
            agent_id: 智能体ID
            
        Returns:
            np.ndarray: 形状为(32, 32, 3, 5)的观测地图
        """
        obs_map = np.zeros((32, 32, 3, 5), dtype=np.float32)

        # 获取当前智能体信息
        if agent_id >= len(self.agents_info):
            return obs_map

        current_agent = self.agents_info[agent_id]

        # 将地图范围映射到32x32网格
        lon_range = self.map_lon_max - self.map_lon_min
        lat_range = self.map_lat_max - self.map_lat_min
        alt_range = self.altitude_max - self.altitude_min

        # 使用新的威胁区计算器计算威胁强度
        # 不再需要预计算威胁区域边界，而是动态计算每个网格点的威胁强度

        # 填充观测地图
        for x in range(32):
            for y in range(32):
                for h in range(3):  # 3个高度层
                    # 特征0: 敌方飞机数量
                    enemy_count = 0
                    for enemy in self.enemies_info:
                        enemy_x = int((enemy.Longtitude - self.map_lon_min) / lon_range * 32)
                        enemy_y = int((enemy.Latitude - self.map_lat_min) / lat_range * 32)
                        enemy_h = int((enemy.TargetAlt - self.altitude_min) / alt_range * 3)
                        if (0 <= enemy_x < 32 and 0 <= enemy_y < 32 and
                                0 <= enemy_h < 3 and enemy_h == h):
                            enemy_count += 1
                    obs_map[x, y, h, 0] = enemy_count

                    # 特征1: 我方飞机数量
                    friendly_count = 0
                    for agent in self.agents_info:
                        agent_x = int((agent.Longtitude - self.map_lon_min) / lon_range * 32)
                        agent_y = int((agent.Latitude - self.map_lat_min) / lat_range * 32)
                        agent_h = int((agent.Altitude - self.altitude_min) / alt_range * 3)
                        if (0 <= agent_x < 32 and 0 <= agent_y < 32 and
                                0 <= agent_h < 3 and agent_h == h):
                            friendly_count += 1
                    obs_map[x, y, h, 1] = friendly_count

                    # 特征2: 威胁区域体积占比（每个小立体空间内半球威胁区的体积占比）
                    threat_volume_ratio = self._calculate_threat_zone_ratio_in_cell(x, y, h)
                    obs_map[x, y, h, 2] = threat_volume_ratio

                    # 特征3: 当前遍历的飞机是否在这个区域内
                    current_agent_x = int((current_agent.Longtitude - self.map_lon_min) / lon_range * 32)
                    current_agent_y = int((current_agent.Latitude - self.map_lat_min) / lat_range * 32)
                    current_agent_h = int((current_agent.Altitude - self.altitude_min) / alt_range * 3)
                    obs_map[x, y, h, 3] = 1.0 if (
                            x == current_agent_x and y == current_agent_y and h == current_agent_h) else 0.0

                    # 特征4: 存在的导弹数量
                    missile_count = 0
                    for agent in self.agents_info:
                        missile_count += agent.MissileNowNum
                    obs_map[x, y, h, 4] = missile_count

        return obs_map

    def _create_36d_obs_vector(self, agent_id: int) -> np.ndarray:
        """
        创建36维观测向量
        
        Args:
            agent_id: 智能体ID
            
        Returns:
            np.ndarray: 形状为(36,)的观测向量
        """
        # 初始化36维观测向量
        obs_vector = np.zeros(36, dtype=np.float32)

        # 获取当前智能体信息
        if agent_id >= len(self.agents_info):
            return obs_vector

        current_agent = self.agents_info[agent_id]

        # 前12维：当前智能体的状态信息 (0-11)
        obs_vector[0] = current_agent.Longtitude  # 经度
        obs_vector[1] = current_agent.Latitude   # 纬度
        obs_vector[2] = current_agent.Altitude   # 高度
        obs_vector[3] = current_agent.Yaw        # 航向角
        obs_vector[4] = current_agent.Pitch      # 俯仰角
        obs_vector[5] = current_agent.Roll       # 滚转角
        obs_vector[6] = current_agent.V_N        # 北向速度
        obs_vector[7] = current_agent.V_E        # 东向速度
        obs_vector[8] = current_agent.V_U        # 天向速度
        obs_vector[9] = current_agent.Mach_M     # 马赫数
        obs_vector[10] = current_agent.MissileNowNum  # 导弹数量
        obs_vector[11] = 1.0 if current_agent.isRadiateOn else 0.0  # 雷达状态

        # 中间12维：敌方信息 (12-23)
        enemy_idx = 0
        for enemy in self.enemies_info[:4]:  # 最多4个敌方信息
            start_idx = 12 + enemy_idx * 3
            obs_vector[start_idx] = enemy.TargetDis      # 距离
            obs_vector[start_idx + 1] = enemy.TargetYaw  # 方位角
            obs_vector[start_idx + 2] = enemy.TargetPitch  # 俯仰角
            enemy_idx += 1

        # 后12维：友方信息 (24-35)
        friendly_idx = 0
        for agent in self.agents_info[:4]:  # 最多4个友方信息
            if agent.DroneID != current_agent.DroneID:  # 排除自己
                start_idx = 24 + friendly_idx * 3
                obs_vector[start_idx] = agent.Longtitude - current_agent.Longtitude  # 相对经度
                obs_vector[start_idx + 1] = agent.Latitude - current_agent.Latitude   # 相对纬度
                obs_vector[start_idx + 2] = agent.Altitude - current_agent.Altitude   # 相对高度
                friendly_idx += 1

        return obs_vector

    def _action_to_control(self, agent_id: int, action: int) -> Dict:
        """
        将动作转换为控制指令
        
        Args:
            agent_id: 智能体ID
            action: 离散动作索引 (0-6)
            
        Returns:
            Dict: 控制指令字典
        """
        # 解析离散动作索引（0-6）
        # 0: 等速平飞
        # 1: 加速
        # 2: 减速
        # 3: 爬升
        # 4: 转弯
        # 5: 俯冲
        # 6: 攻击
        
        # 获取动作类型
        move_type = action if action < 6 else 0  # 如果是攻击动作，使用等速平飞
        attack = action == 6  # 第6个动作是攻击
        
        # 根据动作类型设置移动方向
        if move_type == 0:  # 等速平飞
            x, y, h = 0, 1, 0  # 向北飞行
        elif move_type == 1:  # 加速
            x, y, h = 0, 1, 0  # 向北加速
        elif move_type == 2:  # 减速
            x, y, h = 0, -1, 0  # 向北减速
        elif move_type == 3:  # 爬升
            x, y, h = 0, 0, 1  # 垂直爬升
        elif move_type == 4:  # 转弯
            x, y, h = 1, 0, 0  # 向东转弯
        elif move_type == 5:  # 俯冲
            x, y, h = 0, 0, -1  # 垂直俯冲
        else:
            x, y, h = 0, 1, 0  # 默认向北飞行

        # 构建控制指令
        control_info = {
            'DroneID': agent_id,
            'MoveType': move_type,
            'MoveDirection': [x, y, h],
            'Attack': attack
        }

        return control_info

    def _heuristic_attack(self, agent_info: DroneInfo) -> Optional[int]:
        """
        启发式攻击算法：选择最近的敌机进行攻击
        
        Args:
            agent_info: 智能体信息
            
        Returns:
            Optional[int]: 目标敌机ID，如果没有目标则返回None
        """
        if not agent_info.FoundEnemyList:
            return None

        # 选择距离最近的敌机
        nearest_enemy = min(agent_info.FoundEnemyList, key=lambda x: x.TargetDis)
        return nearest_enemy.EnemyID

    def reset(self):
        """
        重置环境
        
        Returns:
            tuple: (obs, available_actions) 与环境包装器期望的格式匹配
        """
        self.step_count = 0
        self.episode_reward = 0
        self.done = False

        # 重新连接仿真平台
        if not self.connected:
            self._connect_to_simulation()

        # 获取初始观测
        obs_info = self._receive_obs_info()
        self._parse_obs_info(obs_info)

        # 创建初始观测
        obs = self.get_obs()
        available_actions = np.array([self.get_avail_agent_actions(i) for i in range(self.n_agents)])

        return obs, available_actions

    def step(self, actions: np.ndarray):
        """
        执行一步动作
        Args:
            actions: 智能体动作数组
            
        Returns:
            tuple: (obs, rewards, dones, infos, available_actions) 与环境包装器期望的格式匹配
        """
        self.step_count += 1

        # 发送控制指令
        actions = np.array(actions)  # 确保actions是numpy数组
        
        # 处理批次维度
        if len(actions.shape) == 3 and actions.shape[0] == 1:  # 如果是(1, n_agents, 1)
            actions = actions.squeeze(0)  # 变成(n_agents, 1)
        
        # 处理每个智能体的动作
        for i in range(min(len(self.agents_info), len(actions))):
            # 获取当前智能体的动作
            agent_action = actions[i]
            if isinstance(agent_action, np.ndarray):
                agent_action = agent_action.squeeze()  # 移除多余的维度
            
            # 将连续值映射到离散动作（假设动作是logits）
            if isinstance(agent_action, np.ndarray) and agent_action.size > 1:
                action_idx = np.argmax(agent_action)  # 选择最大值对应的索引
            else:
                # 如果是单个值，将其映射到[0,5]范围
                action_idx = int(np.clip((agent_action + 2) * 1.5, 0, 5))
            
            # 执行动作
            control_info = self._action_to_control(i, action_idx)
            self._send_control_command(i, control_info)

        # 接收新的观测信息
        obs_info = self._receive_obs_info()
        self._parse_obs_info(obs_info)

        # 计算奖励
        rewards = self._calculate_rewards()  # (n_agents,)
        rewards = rewards.reshape(1, -1)  # (1, n_agents)

        # 检查是否完成
        self.done = self._check_done()

        # 创建观测
        obs = self.get_obs()  # 返回numpy数组格式，形状为(1, n_agents, *obs_shape)
        
        # 获取可用动作
        available_actions = np.array([self.get_avail_agent_actions(i) for i in range(self.n_agents)])
        available_actions = available_actions.reshape(1, *available_actions.shape)  # (1, n_agents, action_dim)

        # 更新总奖励
        self.episode_reward += np.sum(rewards)

        # 创建每个智能体的done状态
        dones = np.array([self.done for _ in range(self.n_agents)])  # (n_agents,)
        dones = dones.reshape(1, -1)  # (1, n_agents)

        # 创建信息字典
        infos = [{'episode_reward': self.episode_reward, 'step_count': self.step_count}]

        return obs, rewards, dones, infos, available_actions

    def _parse_obs_info(self, obs_info: Dict):
        """解析观测信息"""
        self.agents_info = []
        self.enemies_info = []

        # 解析智能体信息
        if 'agents' in obs_info:
            for agent_data in obs_info['agents']:
                agent = DroneInfo(agent_data)
                self.agents_info.append(agent)

        # 解析敌机信息
        if 'enemies' in obs_info:
            for enemy_data in obs_info['enemies']:
                enemy = EnemyInfo(enemy_data)
                self.enemies_info.append(enemy)

    def _calculate_rewards(self) -> np.ndarray:
        """计算奖励"""
        rewards = np.zeros(self.n_agents)

        for i, agent in enumerate(self.agents_info):
            # 基础奖励
            reward = 0

            # 存活奖励
            if agent.Altitude > 0:
                reward += 1

            # 击落敌机奖励
            for enemy in self.enemies_info:
                if enemy.TargetDis < 1000:  # 距离小于1km
                    reward += 10

            # 被击落惩罚
            if agent.Altitude <= 0:
                reward -= 50

            rewards[i] = reward

        return rewards

    def _check_done(self) -> bool:
        """检查episode是否结束"""
        # 达到最大步数
        if self.step_count >= self.episode_limit:
            return True

        # 所有智能体被击落
        if all(agent.Altitude <= 0 for agent in self.agents_info):
            return True

        # 所有敌机被击落
        if len(self.enemies_info) == 0:
            return True

        return False

    def get_obs(self):
        """返回所有智能体的观测作为numpy数组
        用于与shared_buffer兼容
        返回形状：(n_rollout_threads, n_agents, *obs_shape)
        """
        agents_obs = [self.get_obs_agent(i) for i in range(self.n_agents)]
        # 将列表转换为numpy数组，形状为(n_agents, *obs_shape)
        obs = np.array(agents_obs)
        # 添加n_rollout_threads维度
        obs = obs.reshape(1, *obs.shape)  # (1, n_agents, *obs_shape)
        return obs

    def get_obs_agent(self, agent_id: int) -> np.ndarray:
        """获取指定智能体的观测"""
        return self._create_36d_obs_vector(agent_id)  # 返回36维观测向量

    def get_avail_agent_actions(self, agent_id: int) -> np.ndarray:
        """获取指定智能体的可用动作"""
        # 创建可用动作掩码 (7个动作)
        avail_actions = np.zeros(7, dtype=np.float32)

        # 所有智能体都允许所有动作
        avail_actions[:] = 1.0

        return avail_actions

    def get_obs_size(self) -> int:
        """获取观测空间大小"""
        return 36  # 36维状态向量

    def get_avail_actions_size(self) -> int:
        """获取可用动作空间大小"""
        return 7  # 7个动作

    def get_total_actions(self) -> int:
        """获取总动作数"""
        return 7  # 7个动作

    def close(self):
        """关闭环境"""
        self._disconnect_from_simulation()

    def seed(self, seed=None):
        """设置随机种子"""
        np.random.seed(seed)
        return [seed]
